import numpy as np


# scripts/gen_data.py 修正部分
def gen_golden_data_simple():
    batch = 8
    channels = 2048
    momentum = 0.9
    epsilon = 1e-5
    total_length = batch * channels  # 正确计算总长度：16384

    # 正确定义tiling结构体（与C++中SyncBNTrainingUpdateCustomTilingData对应）
    tiling = np.zeros((1), dtype=[
        ('totalLength', np.uint32),
        ('tileNum', np.uint32),
        ('tileLength', np.uint32),
        ('lasttileLength', np.uint32),
        ('momentum', np.float32),
        ('epsilon', np.float32),
        ('blockLength', np.uint32),
        ('isEvenCore', np.uint32),
        ('formerNum', np.uint32),
        ('formerLength', np.uint32),
        ('formerTileNum', np.uint32),
        ('formerTileLength', np.uint32),
        ('formerLastTileLength', np.uint32),
        ('tailNum', np.uint32),
        ('tailLength', np.uint32),
        ('tailTileNum', np.uint32),
        ('tailTileLength', np.uint32),
        ('tailLastTileLength', np.uint32),
    ])
    tiling['totalLength'] = total_length  # 关键：设置总长度为16384
    tiling['momentum'] = momentum
    tiling['epsilon'] = epsilon

    # 生成输入数据（保持不变）
    local_mean = np.random.uniform(0, 10, [batch, channels]).astype(np.float32)
    local_var = np.random.uniform(0.1, 5, [batch, channels]).astype(np.float32)
    global_mean = np.random.uniform(0, 10, [batch, channels]).astype(np.float32)
    global_var = np.random.uniform(0.1, 5, [batch, channels]).astype(np.float32)

    # 计算参考结果（保持不变）
    golden_updated_mean = momentum * global_mean + (1 - momentum) * local_mean
    golden_updated_var = momentum * global_var + (1 - momentum) * local_var + epsilon

    # 写入文件（删除错误的tiling覆盖，直接写入结构体）
    tiling.tofile("./input/input_tiling.bin")  # 正确写入包含totalLength的结构体
    local_mean.tofile("./input/input_local_mean.bin")
    local_var.tofile("./input/input_local_var.bin")
    global_mean.tofile("./input/input_global_mean.bin")
    global_var.tofile("./input/input_global_var.bin")
    golden_updated_mean.tofile("./output/golden_updated_mean.bin")
    golden_updated_var.tofile("./output/golden_updated_var.bin")

if __name__ == "__main__":
    gen_golden_data_simple()
